//!
/**
 ***********************************************************************************************************
 * <RBD3D11EffectCompiler implementation>
 ***********************************************************************************************************
 * @file RBD3D11EffectCompiler.cpp
 ***********************************************************************************************************/

/************************************************************************************************************
 * Includes
 ************************************************************************************************************/
#include "RBD3D11EffectCompiler.h"

#include <d3dx11.h>
#include <d3dcompiler.h>

#include "RBD3D11Renderer.h"
#include "RBRenderProcess/RBEffect.h"
#include "RBRenderProcess/RBEffectFactory.h"
#include "RBRenderProcess/RBEffectTechnique.h"
#include "RBRenderProcess/RBEffectPass.h"
#include "RBRenderProcess/RBEffectVariable.h"
#include "RBRenderProcess/RBShader.h"

#include "RBMain/RBFile.h"

#include <FXParser.h>

/************************************************************************************************************/


/***********************************************************************************************************
 * 
 ***********************************************************************************************************/
RBD3D11EffectCompiler::RBD3D11EffectCompiler( RBD3D11Renderer* pRenderer)
: RBEffectCompiler()
, m_pRenderer(pRenderer)
{

}

void CompileShader(RBShader* pShader, ID3D11Device* pD3DDevice, const RBString& text, const String& entryPoint, const String& profile, uint32 uiShader, ID3D10Blob** ppBlob = NULL)
{
    ID3D10Blob* blob = NULL;
    ID3D10Blob* error = NULL;
    D3D10_SHADER_MACRO macros[] = {NULL};
    // Shader was not compiled with this profile: add it
    HRESULT hr = D3DCompile(text.c_str(), text.size(), NULL, macros, NULL, 
                            entryPoint.c_str(), profile.c_str(), D3D10_SHADER_ENABLE_BACKWARDS_COMPATIBILITY, 0, &blob, &error);

    if (SUCCEEDED(hr))
    {
        IUnknown* pNative = NULL;
        switch (uiShader)
        {
        case ShaderType::VertexShader:   hr = pD3DDevice->CreateVertexShader   (blob->GetBufferPointer(), blob->GetBufferSize(), NULL, (ID3D11VertexShader**)&pNative);   break;
        case ShaderType::PixelShader:    hr = pD3DDevice->CreatePixelShader    (blob->GetBufferPointer(), blob->GetBufferSize(), NULL, (ID3D11PixelShader**)&pNative);    break;
        case ShaderType::ComputeShader:  hr = pD3DDevice->CreateComputeShader  (blob->GetBufferPointer(), blob->GetBufferSize(), NULL, (ID3D11ComputeShader**)&pNative);  break;
        case ShaderType::GeometryShader: hr = pD3DDevice->CreateGeometryShader (blob->GetBufferPointer(), blob->GetBufferSize(), NULL, (ID3D11GeometryShader**)&pNative); break;
        default:     hr = E_FAIL; break;
        }

        if (SUCCEEDED(hr))
        {
            pShader->Native = pNative;
            if (ppBlob != NULL)
            {
                *ppBlob = blob;
            }
            else
            {
                RBCommon::SafeRelease(blob);
            }
        }
    }

    RBCommon::SafeRelease(error);

}

void IntrospectCommandBuffers(D3D11_SHADER_DESC d3dShaderDesc, ID3D11ShaderReflection* pReflector,
                              ID3D11Device* pD3DDevice, RBShader* pShader,
                              std::vector<RBEffectCBVariable>* pEffectCBVariables, stdext::hash_map<RBString, uint32>* pEffectCBVariablesNames, uint16& uiEffectBufferOffset)
{

    for (uint32 uiCB=0; uiCB < d3dShaderDesc.ConstantBuffers; ++uiCB)
    {
        ID3D11ShaderReflectionConstantBuffer* pD3DCBReflection = pReflector->GetConstantBufferByIndex(uiCB);
        D3D11_SHADER_BUFFER_DESC d3dCBDesc;
        pD3DCBReflection->GetDesc(&d3dCBDesc);
        D3D11_BUFFER_DESC d3dBufferDesc = 
        {
            d3dCBDesc.Size,
            D3D11_USAGE_DYNAMIC,
            D3D11_BIND_CONSTANT_BUFFER,
            D3D11_CPU_ACCESS_WRITE,
            0,
            d3dCBDesc.Size
        };

        RBShaderCommandBuffer shaderCommandBuffer;

        ID3D11Buffer* pD3DCB = 0;
        pD3DDevice->CreateBuffer(&d3dBufferDesc, NULL, &pD3DCB);
        shaderCommandBuffer.Native = pD3DCB;

        for (uint32 uiVar=0; uiVar < d3dCBDesc.Variables; ++uiVar)
        {
            ID3D11ShaderReflectionVariable* pD3DVar = pD3DCBReflection->GetVariableByIndex(uiVar);
            D3D11_SHADER_VARIABLE_DESC d3dVarDesc;
            pD3DVar->GetDesc(&d3dVarDesc);
            RBShaderConstantBufferVariable var;

            RBString varToFind = RBString(d3dVarDesc.Name);

            stdext::hash_map<RBString, uint32>::iterator iter = pEffectCBVariablesNames->find(varToFind);
            // Create a new effect variable if not already present
            if (iter == pEffectCBVariablesNames->end())
            {
                // Find the variable in the uniforms array
                //uint32 uiVar = FXParser::FindByName(uniforms, varToFind);
                //if (uiVar != -1)
                {
                    //const FXParser::UniformDesc& uniformDesc = uniforms[uiVar];

                    uint32 uiShaderSemantic = RBEffect::get_semantic_from_name(d3dVarDesc.Name);

                    // Only float and int are supported for now
                        {
                            RBEffectCBVariable effectVariable = 
                            {
                                d3dVarDesc.Name,
                                uiShaderSemantic,
                                uiEffectBufferOffset,
                                (uint16)d3dVarDesc.Size
                            };
                            uiEffectBufferOffset += (uint16)d3dVarDesc.Size;
                            uiVar = pEffectCBVariables->size();
                            pEffectCBVariables->push_back(effectVariable);
                            iter = pEffectCBVariablesNames->insert(std::make_pair(varToFind, uiVar)).first;
                        }
                }
            }
            // If the variable is now in the effect buffer array, push the references to the shader array
            if (iter != pEffectCBVariablesNames->end())
            {
                const RBEffectCBVariable& effectCBVar = pEffectCBVariables->at(iter->second);
                var.Size                 = (uint16)effectCBVar.Size;
                var.EffectBufferOffset   = effectCBVar.Offset;
                var.ConstantBufferOffset = (uint16)d3dVarDesc.StartOffset;
                shaderCommandBuffer.ConstantBufferVariables.push_back(var);
            }

        }
        pShader->CommandBuffers.push_back(shaderCommandBuffer);
    }
}

template <class TVariable>
uint32 IntrospectRes(std::vector<TVariable>* pEffectVariables, stdext::hash_map<RBString, uint32>* pEffectVariablesNames,
                     const D3D11_SHADER_INPUT_BIND_DESC& d3dResourceDesc)
{
    uint32 uiShaderSemantic = RBEffect::get_semantic_from_name(d3dResourceDesc.Name);

    TVariable effectVariable;
    effectVariable.Semantic = uiShaderSemantic;
    effectVariable.Name = d3dResourceDesc.Name;
    uint32 uiVar = pEffectVariables->size();
    pEffectVariables->push_back(effectVariable);
    pEffectVariablesNames->insert(std::make_pair(d3dResourceDesc.Name, uiVar));
    return uiVar;
}

void IntrospectResources(D3D11_SHADER_DESC d3dShaderDesc, ID3D11ShaderReflection* pReflector, RBShader* pShader,
                         std::vector<RBEffectRVVariable>* pEffectRVVariables, stdext::hash_map<RBString, uint32>* pEffectRVVariablesNames,
                         std::vector<RBEffectSamplerVariable>* pEffectSamplersVariables, stdext::hash_map<RBString, uint32>* pEffectSamplerVariablesNames)
{
    for (uint32 uiResource = 0; uiResource < d3dShaderDesc.BoundResources; ++uiResource)
    {
        D3D11_SHADER_INPUT_BIND_DESC d3dResourceDesc;
        pReflector->GetResourceBindingDesc(uiResource, &d3dResourceDesc);
        if (d3dResourceDesc.Type == D3D_SIT_TEXTURE)
        {
            uint32 idx = IntrospectRes(pEffectRVVariables, pEffectRVVariablesNames, d3dResourceDesc);

            {
                RBShaderResourceViewVariable rscView; 
                rscView.BindPoint = d3dResourceDesc.BindPoint;
                rscView.EffectViewIndex = idx;
                pShader->ResourceViewVariables.push_back(rscView);
            }

        }
        else if (d3dResourceDesc.Type == D3D_SIT_SAMPLER)
        {
            uint32 idx = IntrospectRes(pEffectSamplersVariables, pEffectSamplerVariablesNames, d3dResourceDesc);
            {
                RBShaderSamplerVariable samplerVar;
                samplerVar.BindPoint = d3dResourceDesc.BindPoint;
                samplerVar.EffectViewIndex = idx;
                pShader->SamplerVariables.push_back(samplerVar);
            }

        }
    }
}

void IntrospectShader(ID3D11Device* pD3DDevice, RBShader* pShader, ID3D10Blob* pCodeBlob,
                      std::vector<RBEffectCBVariable>* pEffectCBVariables, stdext::hash_map<RBString, uint32>* pEffectCBVariablesNames, uint16& uiEffectBufferOffset,
                      std::vector<RBEffectRVVariable>* pEffectRVVariables, stdext::hash_map<RBString, uint32>* pEffectRVVariablesNames,
                      std::vector<RBEffectSamplerVariable>* pEffectSamplersVariables, stdext::hash_map<RBString, uint32>* pEffectSamplerVariablesNames)
{
    ID3D11ShaderReflection* pReflector = NULL;
    RBAssert(pCodeBlob!= NULL, "invalid blob !");
    D3DReflect(pCodeBlob->GetBufferPointer(), pCodeBlob->GetBufferSize(), IID_ID3D11ShaderReflection, (void**)&pReflector);
    if (!pReflector)
    {
        return;
    }

    D3D11_SHADER_DESC d3dShaderDesc;
    pReflector->GetDesc(&d3dShaderDesc);

    IntrospectCommandBuffers(d3dShaderDesc, pReflector, pD3DDevice, pShader,  pEffectCBVariables, pEffectCBVariablesNames, uiEffectBufferOffset);
    IntrospectResources(d3dShaderDesc, pReflector, pShader, pEffectRVVariables, pEffectRVVariablesNames, pEffectSamplersVariables, pEffectSamplerVariablesNames);


}

/***********************************************************************************************************
 * 
 ***********************************************************************************************************/
bool RBD3D11EffectCompiler::CompileNative( const RBString& text, RBEffect& kEffect, const FXParser::FXDesc& fxDesc )
{
    ID3D11Device* pD3DDevice = m_pRenderer->GetDevice();
    {
        std::vector<RBEffectCBVariable> vEffectCBVariables;
        std::vector<RBEffectRVVariable> vEffectRVVariables;
        std::vector<RBEffectSamplerVariable> vEffectSamplersVariables;

        stdext::hash_map<RBString, uint32> mEffectCBVariablesNames;
        stdext::hash_map<RBString, uint32> mEffectRVVariablesNames;
        stdext::hash_map<RBString, uint32> mEffectSamplerVariablesNames;

        uint16 uiEffectBufferOffset=0;
        for (uint32 uiTechnique=0; uiTechnique<kEffect.GetTechniqueCount(); ++uiTechnique)
        {
            RBEffectTechnique* pTechnique = kEffect.GetTechnique(uiTechnique);
            for (uint32 uiPass=0; uiPass<pTechnique->GetPassCount(); ++uiPass)
            {
                RBEffectPass* pPass = pTechnique->GetPass(uiPass);
                for (uint32 uiShader=0; uiShader<EST_COUNT; ++uiShader)
                {
                    RBShader* pShader = pPass->GetShader((EShaderType)uiShader);
                    if (pShader == NULL)
                    {
                        continue;
                    }

                    ID3D10Blob* pCodeBlob = NULL;
                    {
                        CompileShader(pShader, pD3DDevice, text, pShader->EntryPoint, pShader->Profile, uiShader, &pCodeBlob);
                        if (pShader != NULL)
                        {
                            IntrospectShader(pD3DDevice, pShader, pCodeBlob,
                                             &vEffectCBVariables,       &mEffectCBVariablesNames, uiEffectBufferOffset,
                                             &vEffectRVVariables,       &mEffectRVVariablesNames,
                                             &vEffectSamplersVariables, &mEffectSamplerVariablesNames );
                        }
                    }

                    if (uiShader==EST_VS)
                    {
                        pPass->CreateIASignature(pCodeBlob->GetBufferPointer(), pCodeBlob->GetBufferSize());
                    }

                    RBCommon::SafeRelease(pCodeBlob);
                }
            }
        }

        kEffect.SetEffectVariables(vEffectCBVariables, vEffectRVVariables, vEffectSamplersVariables);

        return true;
    }

    return false;
}

struct RBD3D11Includer : public ID3D10Include
{
    STDMETHOD(Open)(D3D10_INCLUDE_TYPE IncludeType, LPCSTR pFileName, LPCVOID pParentData, LPCVOID *ppData, UINT *pBytes);
    STDMETHOD(Close)(LPCVOID pData);
    std::vector<RBString> m_vDirectories;
    std::vector<RBString> m_currentFilenames;
};

#pragma warning(push)
#pragma warning(disable: 4100)

HRESULT RBD3D11Includer::Open(D3D10_INCLUDE_TYPE IncludeType, LPCSTR pFileName, LPCVOID pParentData, LPCVOID *ppData, UINT *pBytes)
{
    RBFile* pFile = NULL;

    RBString strNewFilename;


    if (IncludeType == D3D10_INCLUDE_LOCAL)
    {
        strNewFilename = RBStringTools::GetDirName(m_currentFilenames.back()) + "/" + RBString(pFileName);
        pFile = RBFile::Open(strNewFilename.c_str());
    }
    else
    {
        for (uint32 uiDir=0; uiDir < m_vDirectories.size() && (pFile==NULL); ++uiDir)
        {
            strNewFilename = m_vDirectories[uiDir] + "/" + pFileName;
            pFile = RBFile::Open(strNewFilename.c_str());
        }
    }


    if (!pFile)
    {
        return E_FAIL;
    }

    m_currentFilenames.push_back(strNewFilename);

    char* pBuffer = (char*)malloc(pFile->GetSize());
    pFile->Read(pBuffer, pFile->GetSize());
    *ppData = pBuffer;
    *pBytes = pFile->GetSize();

    RBFile::Close(pFile);

    return S_OK;
}

HRESULT RBD3D11Includer::Close(LPCVOID pData)
{
    if (pData)
    {
        free((void*)pData);
    }
    m_currentFilenames.pop_back();
    return S_OK;
}
#pragma warning(pop)

/***********************************************************************************************************
 * 
 ***********************************************************************************************************/
bool RBD3D11EffectCompiler::PreprocessShaderFile( RBString* pShaderText, const RBString& strFilename )
{
    D3D10_SHADER_MACRO d3dDefines[] = {NULL};
    RBD3D11Includer kIncluder;
    kIncluder.m_currentFilenames.push_back(strFilename);
    RBEffectFactory::GetSingleton()->GetDirectories(kIncluder.m_vDirectories);
    
    ID3D10Blob* pShader = NULL;
    ID3D10Blob* pErrors = NULL;
    HRESULT hr = D3DX11PreprocessShaderFromFileA(strFilename.c_str(), d3dDefines, &kIncluder, NULL, &pShader, &pErrors, NULL);
    
    const char* shaderText = (char*)pShader->GetBufferPointer();
    if (FAILED(hr))
    {
        RBCommon::SafeRelease(pShader);
        RBCommon::SafeRelease(pErrors);
        return false;
    }

    *pShaderText = shaderText;

    RBCommon::SafeRelease(pShader);
    RBCommon::SafeRelease(pErrors);

    return true;
}